exp_dir=/home/node27_tmpdata/xlgeng/pachong_10W_data/fairseq_data/hubert_feat/ximalaya_redian_2T
data_dir=/home/node27_tmpdata/xlgeng/pachong_10W_data/fairseq_data/manifest/ximalaya_redian_2T
exp_name=hubert_large_by_xlgeng
all_data=( "train" )
extract_layer=9            # which layer to extract hiddens features
feat=${exp_name}_extract_layer_${extract_layer}
feat_dir=${exp_dir}/feature/${feat}
nj=40
num_gpu=8
km_dir=${exp_dir}/k-means/${feat}
mkdir -p ${km_dir}

stage=14
stop_stage=14
if [ ${stage} -le 13 ] && [ ${stop_stage} -ge 13 ]; then
    echo "stage 13: K-means application started @ `date`"
    for set in "${all_data[@]}"; do
        echo "drop label: $set"
        for ((i = 0; i < $nj; ++i)); do
        {
            result=$((i % num_gpu))
            export CUDA_VISIBLE_DEVICES=$result
            python tools/dump_km_label.py         \
                ${feat_dir}/${set} ${set} ${km_dir}/model.mdl               \
                ${nj} $i ${km_dir}/${set}_label
        } &
        done
        wait
        # merge labels for different shards
        for rank in $(seq 0 $((nj - 1))); do
            cat ${km_dir}/${set}_label/${set}_${rank}_${nj}.km
        done > ${km_dir}/${set}.km
    done
    echo "stage 13: Done @ `date`"
fi
train_set=train
if [ ${stage} -le 14 ] && [ ${stop_stage} -ge 14 ]; then
    echo "stage 14: Generate dict started @ `date`"
        # fairseq-preprocess                              \
        python ../../../fairseq_cli/preprocess.py                 \
            --only-source                               \
            --source-lang km                            \
            --dict-only                                 \
            --trainpref ${km_dir}/${train_set}                 \
            --destdir ${km_dir}                         \
            --workers 10
    echo "stage 14: Done @ `date`"
fi